package com.zwh.common.codec;

import org.springframework.util.StopWatch;

import javax.crypto.*;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.io.*;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.spec.AlgorithmParameterSpec;
import java.util.TimerTask;


/**
 * aes加解密工具类
 * 封装文件加密解密方法
 */
public class MyAesUtils {

    private static final int CACHE_SIZE = 1024;
    private static final String DEFAULT_ENCODING = "UTF-8";

    /**
     * SecretKey 负责保存对称密钥
     */
    private SecretKey secretKey;
    private Cipher cipher;
    /**
     * 算法参数
     */
    private AlgorithmParameterSpec params;

    public MyAesUtils(String algorithm, String key, String iv) throws NoSuchPaddingException, NoSuchAlgorithmException,
            UnsupportedEncodingException {
        this(algorithm, key.getBytes(DEFAULT_ENCODING), iv == null ? null : iv.getBytes(DEFAULT_ENCODING));
    }

    public MyAesUtils(String algorithm, byte[] key, byte[] iv) throws NoSuchPaddingException, NoSuchAlgorithmException {
        this(algorithm, new SecretKeySpec(key, "AES"), iv == null ? null : new IvParameterSpec(iv));
    }

    public MyAesUtils(String algorithm, SecretKey key, AlgorithmParameterSpec paramsSpec) throws NoSuchPaddingException,
            NoSuchAlgorithmException {
        init(algorithm, key);
        this.params = paramsSpec;
    }

    /**
     * 初始化
     *
     * @param algorithm 算法
     * @param key       密钥，如果为{@code null}自动生成一个key
     * @return SymmetricCrypto的子对象，即子对象自身
     */
    public void init(String algorithm, SecretKey key) throws NoSuchPaddingException, NoSuchAlgorithmException {
        if (algorithm == null || key == null) {
            throw new IllegalArgumentException("算法和密钥不能为空");
        }
        this.secretKey = key;
        cipher = Cipher.getInstance(algorithm);
    }

    /**
     * @param mode @link Cipher#ENCRYPT_MODE} 或 {@link Cipher#DECRYPT_MODE}
     * @return
     */
    public void initMode(int mode) throws InvalidAlgorithmParameterException, InvalidKeyException {
        if (params != null) {
            this.cipher.init(mode, this.secretKey, this.params);
        } else {
            this.cipher.init(mode, this.secretKey);
        }
    }

    /**
     * <p>
     * 文件加密
     * </p>
     *
     * @param sourceFilePath 要加密的源文件
     * @param destFilePath   加密后的文件存放路径
     * @throws Exception
     */
    public void encryptFile(String sourceFilePath, String destFilePath) throws Exception {
        File sourceFile = new File(sourceFilePath);
        File destFile = new File(destFilePath);
        if (!sourceFile.exists() && !sourceFile.isFile()) {
            throw new IllegalArgumentException("源文件不存在!");
        }
        if (!destFile.getParentFile().exists()) {
            destFile.getParentFile().mkdirs();
        }
        destFile.createNewFile();

        InputStream in = new FileInputStream(sourceFile);
        OutputStream out = new FileOutputStream(destFile);

        encrypt(in, out, false);
        in.close();
    }

    /**
     * 加密，针对大数据量，可选结束后是否关闭流
     * CipherOutputStream进行加密数据
     *
     * @param in      要加密的输入流
     * @param out     输出流，可以是文件或网络位置
     * @param isClose 是否关闭流
     */
    public void encrypt(InputStream in, OutputStream out, boolean isClose) throws InvalidKeyException,
            InvalidAlgorithmParameterException {
        initMode(Cipher.ENCRYPT_MODE);
        CipherOutputStream cout = new CipherOutputStream(out, this.cipher);
        try {
            byte[] cache = new byte[CACHE_SIZE];
            int nRead = 0;
            while ((nRead = in.read(cache)) != -1) {
                cout.write(cache, 0, nRead);
                cout.flush();
            }
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            try {
                // CipherOutputStream必须关闭，才能完全写出
                cout.close();
                if (isClose) {
                    in.close();
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }




    /**
     * <p>
     * 文件解密
     * </p>
     *
     * @param sourceFilePath
     * @param destFilePath
     * @throws Exception
     */
    public void decryptFile(String sourceFilePath, String destFilePath) throws Exception {
        File sourceFile = new File(sourceFilePath);
        File destFile = new File(destFilePath);
        if (!sourceFile.exists() && !sourceFile.isFile()) {
            throw new IllegalArgumentException("源文件不存在!");
        }

        if (!destFile.getParentFile().exists()) {
            destFile.getParentFile().mkdirs();
        }
        destFile.createNewFile();

        FileInputStream in = new FileInputStream(sourceFile);
        FileOutputStream out = new FileOutputStream(destFile);

        decrypt(in, out, false);
        in.close();
    }

    /**
     * 解密，针对大数据量，可选结束后是否关闭流
     * CipherInputStream进行解密数据
     *
     * @param encryptData 要解密的输入流
     * @param out         输出流，可以是文件或网络位置
     * @param isClose     是否关闭流
     */
    public void decrypt(InputStream encryptData, OutputStream out, boolean isClose) throws InvalidKeyException,
            InvalidAlgorithmParameterException {
        initMode(Cipher.DECRYPT_MODE);
        CipherInputStream cin = new CipherInputStream(encryptData, this.cipher);
        try {
            byte[] cache = new byte[CACHE_SIZE];
            int nRead = 0;
            while ((nRead = cin.read(cache)) != -1) {
                out.write(cache, 0, nRead);
                out.flush();
            }
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            try {
                // OutputStream必须关闭，才能完全写出
                out.close();
                if (isClose) {
                    encryptData.close();
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    public static void main(String[] args) throws Exception {
        String ALGORITHM = "AES/ECB/PKCS5Padding";
        MyAesUtils aes = new MyAesUtils(ALGORITHM, "huangxiaoguo1234", null);

        StopWatch clock = new StopWatch();
        clock.start("任务一");
        aes.encryptFile("F:/拳皇97.7z", "F:/11.encrypt.json");

        clock.stop();
        clock.start("任务二");
        aes.decryptFile("F:/11.encrypt.json", "/F:/拳皇97.1.7z");
        clock.stop();
        System.out.println(clock.prettyPrint());
    }
}
